package com.mrx.loader.protocol.gzip;

import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.net.URLConnection;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.jar.*;

/**
 * 自定义协议的 {@link URLConnection}
 */
public class GzipURLConnection extends URLConnection {

    private final ByteArrayInputStream inputStream;

    private final long contentLength;

    private static final String SCHEME = "gzip";

    private static final Logger logger = LoggerFactory.getLogger(GzipURLConnection.class);
    private static final Map<String, byte[]> TAR_GZ_FILE_CACHE = new HashMap<>();
    private static final Map<String, byte[]> JAR_FILE_CACHE = new ConcurrentHashMap<>();

    /**
     * 构造一个 {@link URLConnection} 对象
     *
     * @param url url
     */
    public GzipURLConnection(URL url) throws IOException {
        super(url);
        String urlString = url.toString();
        String[] split = urlString.split("gzip://");
        // 第一段是 tar.gz 文件路径, 第二段是 tar.gz 文件里的路径
        String tarFile = split[0];
        String jarFile = split[1].split("#")[0];
        URL newURL = new URL(tarFile.replaceFirst(SCHEME, "file"));
        try (InputStream urlInputStream = newURL.openStream()) {
            byte[] tarCachedFile = TAR_GZ_FILE_CACHE.computeIfAbsent(tarFile, key ->
                    loadTar(newURL, urlInputStream, key)
            );
            byte[] jarCachedFile = JAR_FILE_CACHE.computeIfAbsent(jarFile, key ->
                    loadJar(new ByteArrayInputStream(tarCachedFile), key)
            );
            if (jarCachedFile == null) {
                logger.warn("{} not loaded", urlString);
            }
            contentLength = Objects.requireNonNull(jarCachedFile).length;
            inputStream = new ByteArrayInputStream(jarCachedFile);
        }
        connected = false;
    }

    /**
     * 加载 tar 文件
     *
     * @param newURL         tar 文件 url
     * @param urlInputStream tar 文件 inputStream
     * @param tarFile        tar 文件缓存 key
     * @return 读取到的 byte[]
     */
    private byte[] loadTar(URL newURL, InputStream urlInputStream, String tarFile) {
        try {
            logger.debug("loadTar: {} -> {}", newURL, tarFile);
            connected = true;
            return urlInputStream.readAllBytes();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 加载 jar 文件
     *
     * @param inputStream tar 文件的 inputStream
     * @param jarFile     要加载的 jar 文件名
     * @return 读取到的 byte[]
     */
    private byte[] loadJar(ByteArrayInputStream inputStream, String jarFile) {
        byte[] jarCachedFile = null;
        try {
            TarArchiveInputStream tarInputStream = new TarArchiveInputStream(new GzipCompressorInputStream(inputStream));
            for (TarArchiveEntry entry; (entry = tarInputStream.getNextTarEntry()) != null; ) {
                if (!entry.isDirectory() && jarFile.startsWith(entry.getName())) {
                    logger.debug("loadJar: {}", jarFile);
                    jarCachedFile = tarInputStream.readAllBytes();
                    JarInputStream jarInputStream = new JarInputStream(new ByteArrayInputStream(jarCachedFile));
                    Manifest manifest = jarInputStream.getManifest();
                    // 去掉 manifest 中的 Class-Path, 不然 jarLoader 会加载额外的, 错误的 url
                    if (manifest != null) {
                        Attributes mainAttributes = manifest.getMainAttributes();
                        mainAttributes.remove(Attributes.Name.CLASS_PATH);
                        ByteArrayOutputStream out = new ByteArrayOutputStream();
                        JarOutputStream jarOutputStream = new JarOutputStream(out);
                        for (JarEntry e; (e = jarInputStream.getNextJarEntry()) != null; ) {
                            jarOutputStream.putNextEntry(e);
                            jarOutputStream.write(jarInputStream.readAllBytes());
                            jarOutputStream.closeEntry();
                        }
                        jarOutputStream.close();
                        jarOutputStream.flush();
                        jarCachedFile = out.toByteArray();
                    }
                    connected = true;
                }
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return jarCachedFile;
    }

    @Override
    public void connect() {

    }

    @Override
    public InputStream getInputStream() {
        return inputStream;
    }

    @Override
    public long getContentLengthLong() {
        return contentLength;
    }
}
